from PIL import Image
import numpy as np
import os
            
class FaceDataset:
    def __init__(self, path, start = 0):
        self._path = path#母文件夹
        self._start = start
        self._a = self._start
    def pathlist(self):
        lis=[]
        for dirpath,dirnames,filenames in os.walk(self._path):#filenames返回所有子文件路径
            for filename in filenames:
                image_path=os.path.join(dirpath,filename)
                lis.append(image_path)
        return lis
    def image_generator(self):
        for i in self.pathlist():
            img2 = Image.open(i)
            #转换成np.ndarray格式
            img2 = np.array(img2)
            #print(type(img2))
            yield img2 
        return "done"
    def __iter__(self):
        return self
    def __next__(self):
        if self._a <= len(self.pathlist()):
            x=Image.open(self.pathlist()[self._a])
            self._a += 1
            return x
        else:
            raise StopIteration('Traverse images is over!')
    def __len__(self):
        return len(self.pathlist())
    def __getitem__(self,i):
        return np.array(Image.open(self.pathlist()[i]))

face = FaceDataset(r"C:\Users\User\Desktop\code\pycode\lab9\original\\",start = 0)
f = face.image_generator()

#print(next(f))
#print(next(face))
for i in range(face.__len__()):#生成器的测试
    print(next(f))
for j in face:#迭代器的测试
     print(j)
#print(face.__getitem__(0))